/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <set>
#include "parcsr_sor_preconditioner.h"

namespace YHAMG
{

ParCSR_SOR_Preconditioner::ParCSR_SOR_Preconditioner()
	: interior(0),
	label(0),
	num_neighbors(0),
	neighbor_labels(0),
	neighbor_ranks(0),
	recv_ptr(0),
	send_ptr(0),
	send_list(0),
	send_buffer(0),
	Type(2),
	Factor(1.0)
{
}

ParCSR_SOR_Preconditioner::ParCSR_SOR_Preconditioner(int type, int factor)
	: interior(0),
	label(0),
	num_neighbors(0),
	neighbor_labels(0),
	neighbor_ranks(0),
	recv_ptr(0),
	send_ptr(0),
	send_list(0),
	send_buffer(0),
	Type(type),
	Factor(factor)
{
}

ParCSR_SOR_Preconditioner::~ParCSR_SOR_Preconditioner()
{
	if (interior) delete[] interior;
	if (neighbor_labels) delete[] neighbor_labels;
	if (recv_ptr) delete[] recv_ptr;
	if (send_ptr) delete[] send_ptr;
	if (send_list) delete[] send_list;
	if (send_buffer) delete[] send_buffer;
}

void ParCSR_SOR_Preconditioner::Free()
{
	D_Loc.Free();
	D_Recip_Loc.Free();
	L_LocLoc.Free();
	L_LocExt.Free();
	U_LocLoc.Free();
	U_LocExt.Free();
	x_Ext.Free();
	if (interior) delete[] interior;
	if (neighbor_labels) delete[] neighbor_labels;
	if (recv_ptr) delete[] recv_ptr;
	if (send_ptr) delete[] send_ptr;
	if (send_list) delete[] send_list;
	if (send_buffer) delete[] send_buffer;
	comm = 0;
	label = 0;
	num_neighbors = 0;
	interior = 0;
	neighbor_labels = 0;
	recv_ptr = 0;
	send_ptr = 0;
	send_list = 0;
	send_buffer = 0;
}

#define MPI_TAG 550

void ParCSR_SOR_Preconditioner::Setup(const ParCSR_Matrix& A, int REUSE)
{
	int n = A.LocLoc.size[0];
	int Ext_size = A.LocExt.size[1];

	int* A_LocLoc_rowptr = A.LocLoc.rowptr;
	int* A_LocLoc_colind = A.LocLoc.colind;
	double* A_LocLoc_values = A.LocLoc.values;
	int* A_LocExt_rowptr = A.LocExt.rowptr;
	int* A_LocExt_colind = A.LocExt.colind;
	double* A_LocExt_values = A.LocExt.values;

	int* Ext_map = A.Ext_map;

	if (REUSE == 0)
	{
		Free();

		comm = A.comm;

		int rank;
		MPI_Comm_rank(comm, &rank);

		num_neighbors = A.num_neighbors;
		neighbor_ranks = new int[num_neighbors];
		recv_ptr = new int[num_neighbors + 1];
		send_ptr = new int[num_neighbors + 1];
		send_list = new int[A.send_ptr[num_neighbors]];
		send_buffer = new double[A.send_ptr[num_neighbors]];

		for (int r = 0; r < num_neighbors; ++r)
			neighbor_ranks[r] = A.neighbor_ranks[r];
		for (int r = 0; r <= num_neighbors; ++r)
			recv_ptr[r] = A.Ext_ptr[r];
		for (int r = 0; r <= num_neighbors; ++r)
			send_ptr[r] = A.send_ptr[r];
		for (int i = 0; i < A.send_ptr[num_neighbors]; ++i)
			send_list[i] = A.send_list[i];

		interior = new bool[n];
		for (int i = 0; i < n; ++i)
			interior[i] = A_LocExt_rowptr[i + 1] == A_LocExt_rowptr[i];
		for (int i = 0; i < send_ptr[num_neighbors]; ++i)
			interior[send_list[i]] = 0;

		MPI_Request* recv_requests = new MPI_Request[num_neighbors];
		MPI_Request* send_requests = new MPI_Request[num_neighbors];
		MPI_Status status;

		label = 0;
		neighbor_labels = new int[num_neighbors];

		std::set<int> set_neighbor_labels;

		for (int r = 0; r < num_neighbors; ++r)
			MPI_Irecv(neighbor_labels + r, 1, MPI_INT, neighbor_ranks[r], MPI_TAG, comm, recv_requests + r);

		for (int r = 0; r < num_neighbors; ++r)
		{
			if (rank > neighbor_ranks[r])
			{
				MPI_Wait(recv_requests + r, &status);
				set_neighbor_labels.insert(neighbor_labels[r]);
			}
		}

		for (std::set<int>::iterator iter = set_neighbor_labels.begin(); iter != set_neighbor_labels.end(); ++iter, ++label)
			if (*iter != label) break;

		for (int r = 0; r < num_neighbors; ++r)
			MPI_Send(&label, 1, MPI_INT, neighbor_ranks[r], MPI_TAG, comm);

		for (int r = 0; r < num_neighbors; ++r)
			if (rank < neighbor_ranks[r])
				MPI_Wait(recv_requests + r, &status);

		delete[] recv_requests;
		delete[] send_requests;

		int* Ext_from = new int[Ext_size];
		for (int r = 0; r < num_neighbors; ++r)
			for (int i = recv_ptr[r]; i < recv_ptr[r + 1]; ++i)
				Ext_from[i] = r;

		D_Loc.Resize(n);
		D_Recip_Loc.Resize(n);
		
		double* D_Loc_values = D_Loc.values;
		double* D_Recip_Loc_values = D_Recip_Loc.values;

		int* L_LocLoc_rowptr = new int[n + 1];
		int* U_LocLoc_rowptr = new int[n + 1];
		int* L_LocExt_rowptr = new int[n + 1];
		int* U_LocExt_rowptr = new int[n + 1];

		L_LocLoc_rowptr[0] = 0;
		U_LocLoc_rowptr[0] = 0;
		L_LocExt_rowptr[0] = 0;
		U_LocExt_rowptr[0] = 0;

		for (int i = 0, cnt_l_loc = 0, cnt_u_loc = 0, cnt_l_ext = 0, cnt_u_ext = 0; i < n; L_LocLoc_rowptr[++i] = cnt_l_loc, U_LocLoc_rowptr[i] = cnt_u_loc, L_LocExt_rowptr[i] = cnt_l_ext, U_LocExt_rowptr[i] = cnt_u_ext)
		{
			if (interior[i])
			{
				for (int j = A_LocLoc_rowptr[i]; j < A_LocLoc_rowptr[i + 1]; ++j)
				{
					int jcol = A_LocLoc_colind[j];

					if (jcol < i && interior[jcol])
						++cnt_l_loc;
					else if (jcol != i)
						++cnt_u_loc;
				}
			}
			else
			{
				for (int j = A_LocLoc_rowptr[i]; j < A_LocLoc_rowptr[i + 1]; ++j)
				{
					int jcol = A_LocLoc_colind[j];

					if (jcol < i || interior[jcol])
						++cnt_l_loc;
					else if (jcol != i)
						++cnt_u_loc;
				}

				for (int j = A_LocExt_rowptr[i]; j < A_LocExt_rowptr[i + 1]; ++j)
				{
					int jcol = A_LocExt_colind[j];
					int jfrom = Ext_from[jcol];

					if (label > neighbor_labels[jfrom])
						++cnt_l_ext;
					else
						++cnt_u_ext;
				}
			}
		}

		int* L_LocLoc_colind = new int[L_LocLoc_rowptr[n]];
		double* L_LocLoc_values = new double[L_LocLoc_rowptr[n]];
		int* U_LocLoc_colind = new int[U_LocLoc_rowptr[n]];
		double* U_LocLoc_values = new double[U_LocLoc_rowptr[n]];
		int* L_LocExt_colind = new int[L_LocExt_rowptr[n]];
		double* L_LocExt_values = new double[L_LocExt_rowptr[n]];
		int* U_LocExt_colind = new int[U_LocExt_rowptr[n]];
		double* U_LocExt_values = new double[U_LocExt_rowptr[n]];

		for (int i = 0; i < n; ++i)
		{
			if (interior[i])
			{
				for (int j = A_LocLoc_rowptr[i], k = L_LocLoc_rowptr[i], r = U_LocLoc_rowptr[i]; j < A_LocLoc_rowptr[i + 1]; ++j)
				{
					int jcol = A_LocLoc_colind[j];
					double jval = A_LocLoc_values[j];

					if (jcol < i && interior[jcol])
					{
						L_LocLoc_colind[k] = jcol;
						L_LocLoc_values[k++] = jval;
					}
					else if (jcol != i)
					{
						U_LocLoc_colind[r] = jcol;
						U_LocLoc_values[r++] = jval;
					}
					else
					{
						D_Loc_values[i] = jval;
						D_Recip_Loc_values[i] = 1.0 / jval;
					}
				}
			}
			else
			{
				for (int j = A_LocLoc_rowptr[i], k = L_LocLoc_rowptr[i], r = U_LocLoc_rowptr[i]; j < A_LocLoc_rowptr[i + 1]; ++j)
				{
					int jcol = A_LocLoc_colind[j];
					double jval = A_LocLoc_values[j];

					if (jcol < i || interior[jcol])
					{
						L_LocLoc_colind[k] = jcol;
						L_LocLoc_values[k++] = jval;
					}
					else if (jcol != i)
					{
						U_LocLoc_colind[r] = jcol;
						U_LocLoc_values[r++] = jval;
					}
					else
					{
						D_Loc_values[i] = jval;
						D_Recip_Loc_values[i] = 1.0 / jval;
					}
				}

				for (int j = A_LocExt_rowptr[i], k = L_LocExt_rowptr[i], r = U_LocExt_rowptr[i]; j < A_LocExt_rowptr[i + 1]; ++j)
				{
					int jcol = A_LocExt_colind[j];
					double jval = A_LocExt_values[j];
					int jfrom = Ext_from[jcol];

					if (label > neighbor_labels[jfrom])
					{
						L_LocExt_colind[k] = jcol;
						L_LocExt_values[k++] = jval;
					}
					else
					{
						U_LocExt_colind[r] = jcol;
						U_LocExt_values[r++] = jval;
					}
				}
			}
		}

		delete[] Ext_from;

		L_LocLoc.size[0] = n;
		L_LocLoc.size[1] = n;
		L_LocLoc.rowptr = L_LocLoc_rowptr;
		L_LocLoc.colind = L_LocLoc_colind;
		L_LocLoc.values = L_LocLoc_values;

		L_LocExt.size[0] = n;
		L_LocExt.size[1] = Ext_size;
		L_LocExt.rowptr = L_LocExt_rowptr;
		L_LocExt.colind = L_LocExt_colind;
		L_LocExt.values = L_LocExt_values;

		U_LocLoc.size[0] = n;
		U_LocLoc.size[1] = n;
		U_LocLoc.rowptr = U_LocLoc_rowptr;
		U_LocLoc.colind = U_LocLoc_colind;
		U_LocLoc.values = U_LocLoc_values;

		U_LocExt.size[0] = n;
		U_LocExt.size[1] = Ext_size;
		U_LocExt.rowptr = U_LocExt_rowptr;
		U_LocExt.colind = U_LocExt_colind;
		U_LocExt.values = U_LocExt_values;

		x_Ext.Resize(Ext_size);
	}
	else
	{
		int* L_LocLoc_rowptr = L_LocLoc.rowptr;
		int* L_LocLoc_colind = L_LocLoc.colind;
		double* L_LocLoc_values = L_LocLoc.values;
		int* U_LocLoc_rowptr = U_LocLoc.rowptr;
		int* U_LocLoc_colind = U_LocLoc.colind;
		double* U_LocLoc_values = U_LocLoc.values;
		int* L_LocExt_rowptr = L_LocExt.rowptr;
		int* L_LocExt_colind = L_LocExt.colind;
		double* L_LocExt_values = L_LocExt.values;
		int* U_LocExt_rowptr = U_LocExt.rowptr;
		int* U_LocExt_colind = U_LocExt.colind;
		double* U_LocExt_values = U_LocExt.values;
		double* D_Loc_values = D_Loc.values;
		double* D_Recip_Loc_values = D_Recip_Loc.values;

		int* w = new int[n];
		int* v = new int[Ext_size];

		for (int i = 0; i < n; ++i)
			w[i] = -1;
		for (int i = 0; i < Ext_size; ++i)
			v[i] = -1;

		for (int i = 0; i < n; ++i)
		{
			int j0 = A_LocLoc_rowptr[i];
			
			for (int j = A_LocLoc_rowptr[i]; j < A_LocLoc_rowptr[i + 1]; ++j)
				w[A_LocLoc_colind[j]] = j;

			if (w[i] >= j0)
			{
				D_Loc_values[i] = A_LocLoc_values[w[i]];
				D_Recip_Loc_values[i] = 1.0 / A_LocLoc_values[w[i]];
			}
			
			for (int j = L_LocLoc_rowptr[i]; j < L_LocLoc_rowptr[i + 1]; ++j)
				if (w[L_LocLoc_colind[j]] >= j0) L_LocLoc_values[j] = A_LocLoc_values[w[L_LocLoc_colind[j]]];
			for (int j = U_LocLoc_rowptr[i]; j < U_LocLoc_rowptr[i + 1]; ++j)
				if (w[U_LocLoc_colind[j]] >= j0) U_LocLoc_values[j] = A_LocLoc_values[w[U_LocLoc_colind[j]]];
				
			j0 = A_LocExt_rowptr[i];

			for (int j = A_LocExt_rowptr[i]; j < A_LocExt_rowptr[i + 1]; ++j)
				v[A_LocExt_colind[j]] = j;

			for (int j = L_LocExt_rowptr[i]; j < L_LocExt_rowptr[i + 1]; ++j)
				if (v[L_LocExt_colind[j]] >= j0) L_LocExt_values[j] = A_LocExt_values[v[L_LocExt_colind[j]]];
			for (int j = U_LocExt_rowptr[i]; j < U_LocExt_rowptr[i + 1]; ++j)
				if (v[U_LocExt_colind[j]] >= j0) U_LocExt_values[j] = A_LocExt_values[v[U_LocExt_colind[j]]];
		}

		delete[] w;
		delete[] v;
	}
}

int ParCSR_SOR_Preconditioner::InSize() const
{
	return D_Loc.size;
}

int ParCSR_SOR_Preconditioner::OutSize() const
{
	return D_Loc.size;
}

void ParCSR_SOR_Preconditioner::Apply(const Par_Vector& b, const Par_Vector& x) const
{
	int type = Type;
	double omega = Factor;

	int n = D_Loc.size;
	double* D_Loc_values = D_Loc.values;
	double* D_Recip_Loc_values = D_Recip_Loc.values;
	int* L_LocLoc_rowptr = L_LocLoc.rowptr;
	int* L_LocLoc_colind = L_LocLoc.colind;
	double* L_LocLoc_values = L_LocLoc.values;
	int* L_LocExt_rowptr = L_LocExt.rowptr;
	int* L_LocExt_colind = L_LocExt.colind;
	double* L_LocExt_values = L_LocExt.values;
	int* U_LocLoc_rowptr = U_LocLoc.rowptr;
	int* U_LocLoc_colind = U_LocLoc.colind;
	double* U_LocLoc_values = U_LocLoc.values;
	int* U_LocExt_rowptr = U_LocExt.rowptr;
	int* U_LocExt_colind = U_LocExt.colind;
	double* U_LocExt_values = U_LocExt.values;

	double* bv_Loc = b.Loc.values;
	double* xv_Loc = x.Loc.values;
	double* xv_Ext = x_Ext.values;

	if (n == 0) return;

	for (int i = 0; i < n; ++i)
		xv_Loc[i] = bv_Loc[i];

	MPI_Request* recv_requests = new MPI_Request[num_neighbors];
	MPI_Request* send_requests = new MPI_Request[num_neighbors];
	MPI_Status status;

	if (type == 0 || type == 2)
	{
		for (int i = 0; i < n; ++i)
		{
			if (!interior[i]) continue;
			double temp = 0.0;
			for (int j = L_LocLoc_rowptr[i]; j < L_LocLoc_rowptr[i + 1]; ++j)
				temp -= xv_Loc[L_LocLoc_colind[j]] * L_LocLoc_values[j];
			xv_Loc[i] = (xv_Loc[i] + omega * temp) * D_Recip_Loc_values[i];
		}

		for (int r = 0; r < num_neighbors; ++r)
			if (label > neighbor_labels[r] && recv_ptr[r + 1] > recv_ptr[r])
				MPI_Irecv(xv_Ext + recv_ptr[r], recv_ptr[r + 1] - recv_ptr[r], MPI_DOUBLE, neighbor_ranks[r], MPI_TAG, comm, recv_requests + r);

		for (int r = 0; r < num_neighbors; ++r)
			if (label > neighbor_labels[r] && recv_ptr[r + 1] > recv_ptr[r])
				MPI_Wait(recv_requests + r, &status);

		for (int i = 0; i < n; ++i)
		{
			if (interior[i]) continue;
			double temp = 0.0;
			for (int j = L_LocExt_rowptr[i]; j < L_LocExt_rowptr[i + 1]; ++j)
				temp -= xv_Ext[L_LocExt_colind[j]] * L_LocExt_values[j];
			for (int j = L_LocLoc_rowptr[i]; j < L_LocLoc_rowptr[i + 1]; ++j)
				temp -= xv_Loc[L_LocLoc_colind[j]] * L_LocLoc_values[j];
			xv_Loc[i] = (xv_Loc[i] + omega * temp) * D_Recip_Loc_values[i];
		}

		for (int r = 0; r < num_neighbors; ++r)
		{
			if (label < neighbor_labels[r] && send_ptr[r + 1] > send_ptr[r])
			{
				for (int i = send_ptr[r]; i < send_ptr[r + 1]; ++i)
					send_buffer[i] = xv_Loc[send_list[i]];
				MPI_Isend(send_buffer + send_ptr[r], send_ptr[r + 1] - send_ptr[r], MPI_DOUBLE, neighbor_ranks[r], MPI_TAG, comm, send_requests + r);
			}
		}

		for (int r = 0; r < num_neighbors; ++r)
			if (label < neighbor_labels[r] && send_ptr[r + 1] > send_ptr[r])
				MPI_Wait(send_requests + r, &status);
	}

	if (type == 2)
	{
		for (int i = 0; i < n; ++i)
			xv_Loc[i] *= D_Loc_values[i];
	}

	if (type == 1 || type == 2)
	{

		for (int r = 0; r < num_neighbors; ++r)
			if (label < neighbor_labels[r] && recv_ptr[r + 1] > recv_ptr[r])
				MPI_Irecv(xv_Ext + recv_ptr[r], recv_ptr[r + 1] - recv_ptr[r], MPI_DOUBLE, neighbor_ranks[r], MPI_TAG, comm, recv_requests + r);

		for (int r = 0; r < num_neighbors; ++r)
			if (label < neighbor_labels[r] && recv_ptr[r + 1] > recv_ptr[r])
				MPI_Wait(recv_requests + r, &status);

		for (int i = n - 1; i >= 0; --i)
		{
			if (interior[i]) continue;
			double temp = 0.0;
			for (int j = U_LocExt_rowptr[i + 1] - 1; j >= U_LocExt_rowptr[i]; --j)
				temp -= xv_Ext[U_LocExt_colind[j]] * U_LocExt_values[j];
			for (int j = U_LocLoc_rowptr[i + 1] - 1; j >= U_LocLoc_rowptr[i]; --j)
				temp -= xv_Loc[U_LocLoc_colind[j]] * U_LocLoc_values[j];
			xv_Loc[i] = (xv_Loc[i] + omega * temp) * D_Recip_Loc_values[i];
		}

		for (int r = 0; r < num_neighbors; ++r)
		{
			if (label > neighbor_labels[r] && send_ptr[r + 1] > send_ptr[r])
			{
				for (int i = send_ptr[r]; i < send_ptr[r + 1]; ++i)
					send_buffer[i] = xv_Loc[send_list[i]];
				MPI_Isend(send_buffer + send_ptr[r], send_ptr[r + 1] - send_ptr[r], MPI_DOUBLE, neighbor_ranks[r], MPI_TAG, comm, send_requests + r);
			}
		}

		for (int r = 0; r < num_neighbors; ++r)
			if (label > neighbor_labels[r] && send_ptr[r + 1] > send_ptr[r])
				MPI_Wait(send_requests + r, &status);

		for (int i = n - 1; i >= 0; --i)
		{
			if (!interior[i]) continue;
			double temp = 0.0;
			for (int j = U_LocLoc_rowptr[i + 1] - 1; j >= U_LocLoc_rowptr[i]; --j)
				temp -= xv_Loc[U_LocLoc_colind[j]] * U_LocLoc_values[j];
			xv_Loc[i] = (xv_Loc[i] + omega * temp) * D_Recip_Loc_values[i];
		}
	}

	delete[] recv_requests;
	delete[] send_requests;
}

}